import os
import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D, Pool2D, FC
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

from mnist2 import MNIST
from load_data import load_data
#网络结构部分之后的代码，保持不变
place = fluid.CPUPlace()
with fluid.dygraph.guard(place):
    model = MNIST("mnist")
    model.train()
    #调用加载数据的函数，获得MNIST训练数据集
    train_loader = load_data('train')
    # 定义DataLoader对象用于加载Python生成器产生的数据
    data_loader = fluid.io.DataLoader.from_generator(capacity=5, return_list=True)
     # 设置数据生成器
    data_loader.set_batch_generator(train_loader, places=place)
    # 使用SGD优化器，learning_rate设置为0.01
    # optimizer = fluid.optimizer.SGDOptimizer(learning_rate=0.01)
    optimizer=fluid.optimizer.MomentumOptimizer(learning_rate=0.01,momentum=0.9)
    # 训练5轮
    EPOCH_NUM = 5
    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(data_loader):
            #准备数据
            image_data, label_data = data
            image = fluid.dygraph.to_variable(image_data)
            label = fluid.dygraph.to_variable(label_data)
            
            #前向计算的过程
            predict = model(image)
            
            #计算损失，取一个批次样本损失的平均值
            loss = fluid.layers.square_error_cost(predict, label)
            avg_loss = fluid.layers.mean(loss)
            
            #每训练了200批次的数据，打印下当前Loss的情况
            if batch_id % 200 == 0:
                print("epoch: {}, batch: {}, loss is: {}".format(epoch_id, batch_id, avg_loss.numpy()))
            
            #后向传播，更新参数的过程
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()

    #保存模型参数
    fluid.save_dygraph(model.state_dict(), 'mnist')